

#define nThread 128

uint nGroupX = 1;
uint nGroupY = 1;
uint nGroupZ = 1;

Buffer<float> g_srcMat;
RWBuffer<float> g_dstMat;

uint g_RowRange_fromIdx = 0;
uint g_RowRange_ToIdx = 0;
uint g_nRow_srcMat = 1;
uint g_nCol_srcMat = 1;

uint g_nRow_dstMat = 1;
uint g_nCol_dstMat = 1;

Buffer<float> g_src2Mat;
uint g_nRow_src2Mat = 1;
uint g_nCol_src2Mat = 1;

uint g_rowIdx = 0;
uint g_colIdx = 0;

[numthreads( nThread,  1,  1)] //ThreadGroupCount = ceil(nCol*(g_RowRange_ToIdx - g_RowRange_fromIdx + 1) / nThread);
void cs_GetRowRange(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	uint tId_validRange = g_nCol_srcMat * (g_RowRange_ToIdx - g_RowRange_fromIdx + 1);
	if (threadIdx >= tId_validRange)
		return;
	uint nRowGet = (g_RowRange_ToIdx - g_RowRange_fromIdx + 1);
	uint srcRIdx = (int)(threadIdx / nRowGet) + g_RowRange_fromIdx;
	uint srcCIdx = threadIdx % nRowGet;

	uint dstRIdx = srcRIdx - g_RowRange_fromIdx;
	uint dstCIdx = srcCIdx;

	uint addr_src = srcRIdx*g_nCol_srcMat + srcCIdx;
	uint addr_dst = dstRIdx*g_nCol_srcMat + dstCIdx;

	g_dstMat[addr_dst] = g_srcMat[addr_src];

}

uint g_ColRange_fromIdx = 0;
uint g_ColRange_ToIdx = 0;

[numthreads( nThread,  1,  1)] //ThreadGroupCount = ceil(nRow*(g_ColRange_ToIdx - g_ColRange_fromIdx + 1) / nThread);
void cs_GetColRange(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	uint tId_validRange = g_nRow_srcMat * (g_ColRange_ToIdx - g_ColRange_fromIdx + 1);
	if (threadIdx >= tId_validRange)
		return;
	uint nColGet = (g_ColRange_ToIdx - g_ColRange_fromIdx + 1);
	uint srcCIdx = (int)(threadIdx / nColGet) + g_ColRange_fromIdx;
	uint srcRIdx = threadIdx % nColGet;

	uint dstCIdx = srcCIdx - g_ColRange_fromIdx;
	uint dstRIdx = srcRIdx;

	uint addr_src = srcRIdx*g_nCol_srcMat + srcCIdx;
	uint addr_dst = dstRIdx*g_nCol_srcMat + dstCIdx;

	g_dstMat[addr_dst] = g_srcMat[addr_src];
}


int g_diag_shift = 0;
[numthreads( nThread,  1,  1)] //ThreadGroupCount = ceil( min(nRow, nCol) / nThread);
void cs_GetDiag(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	uint tId_validRange = min(g_nRow_srcMat, g_nCol_srcMat);
	if (threadIdx >= tId_validRange)
		return;
	int startCIdx = 0;
	int startRIdx = 0;
	if (g_diag_shift > 0)
	{
		startCIdx += g_diag_shift;
	}
	else if (g_diag_shift < 0)
	{
		startRIdx += abs(g_diag_shift);
	}
	uint srcCIdx = startCIdx + threadIdx;
	uint srcRIdx = startRIdx + threadIdx;
	if (srcCIdx >= g_nCol_srcMat || srcRIdx >= g_nRow_srcMat)
	{
		return;
	}
	uint addr_src = srcRIdx*g_nCol_srcMat + srcCIdx;
	uint addr_dst = threadIdx;

	g_dstMat[addr_dst] = g_srcMat[addr_src];
}


[numthreads( nThread,  1,  1)] 
void cs_SetDiag(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	uint tId_validRange = min(g_nRow_dstMat, g_nCol_dstMat);
	if (threadIdx >= tId_validRange)
		return;
	uint rIdx = threadIdx;
	uint cIdx = threadIdx;
	if (cIdx >= g_nCol_dstMat || rIdx >= g_nRow_dstMat)
	{
		return;
	}
	uint addr_dst = rIdx*g_nCol_dstMat + cIdx;
	uint addr_src = rIdx;

	g_dstMat[addr_dst] = g_srcMat[addr_src];
}


float g_vScalar = 0;
[numthreads( nThread,  1,  1)] //ThreadGroupCount = ceil( min(nRow, nCol) / nThread);
void cs_SetTo(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	uint tId_validRange = g_nRow_srcMat*g_nCol_srcMat;
	if (threadIdx >= tId_validRange)
		return;
	uint rIdx = (int)(threadIdx / g_nCol_srcMat);
	uint cIdx = threadIdx % g_nCol_srcMat;

	uint addr = rIdx*g_nCol_srcMat + cIdx;
	g_dstMat[addr] = g_vScalar;
}


[numthreads( nThread,  1,  1)] //ThreadGroupCount = ceil( nRow*nCol / nThread);
void cs_MaskSetTo(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	uint tId_validRange = g_nRow_srcMat*g_nCol_srcMat;
	if (threadIdx >= tId_validRange)
		return;
	uint rIdx = (int)(threadIdx / g_nCol_dstMat);
	uint cIdx = threadIdx % g_nCol_dstMat;
	uint addr = rIdx*g_nCol_dstMat + cIdx;
	if (g_src2Mat[addr] <= 0)
		return;
	g_dstMat[addr] = g_vScalar;
}

[numthreads( nThread,  1,  1)] //ThreadGroupCount = ceil( nRow*nCol / nThread);
void cs_MaskCopy(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	uint tId_validRange = g_nRow_srcMat*g_nCol_srcMat;
	if (threadIdx >= tId_validRange)
		return;
	uint rIdx = (int)(threadIdx / g_nCol_srcMat);
	uint cIdx = threadIdx % g_nCol_srcMat;
	uint addr = rIdx * g_nCol_srcMat + cIdx;
	if (g_src2Mat[addr] <= 0)
		return;
	g_dstMat[addr] = g_srcMat[addr];
}



[numthreads( nThread,  1,  1)] 
void cs_Identity(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	uint tId_validRange = g_nRow_dstMat*g_nCol_dstMat;
	if (threadIdx >= tId_validRange)
		return;
	uint rIdx = (int)(threadIdx / g_nCol_dstMat);
	uint cIdx = threadIdx % g_nCol_dstMat;
	uint addr = rIdx*g_nCol_dstMat + cIdx;
	g_dstMat[addr] = 0;
	if (rIdx == cIdx)
	{
		g_dstMat[addr] = 1.0;
	}

	return;
}

////////////for matrix mul//////////////
[numthreads( nThread,  1,  1)] //ThreadGroupCount = ceil(g_nRow_src2Mat*g_nCol_src2Mat/ nThread);
void cs_MatrixMul_1RowMul(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	uint tId_validRange = g_nRow_src2Mat*g_nCol_src2Mat;
	if (threadIdx >= tId_validRange)
		return;
	uint cIdx_src2 = (int)(threadIdx / g_nCol_srcMat);
	uint rIdx_src2 = threadIdx % g_nRow_src2Mat;
	uint rIdx_src1 = g_rowIdx;
	uint cIdx_src1 = rIdx_src2;

	uint addr_src1 = rIdx_src1*g_nCol_srcMat + cIdx_src1;
	uint addr_src2 = rIdx_src2*g_nCol_src2Mat + cIdx_src2;

	uint addr_dst = addr_src2;

	g_dstMat[addr_dst] = g_srcMat[addr_src1]*g_src2Mat[addr_src2];

	return;
}

//input: g_srcMat, g_dstMat
//      need log_2(nCol) calls

uint g_nCallStep = 0;
[numthreads( nThread,  1,  1)]    //nThreadGroup = (nCol / 2^g_nCallStep) / nThread 
void cs_MatrixMul_colSum(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	uint nSegSize = pow(2, g_nCallStep);
	uint nSegPerCol = ceil(g_nCol_srcMat/ nSegSize);
	uint rSegIdx = threadIdx / nSegPerCol;
	uint cSegIdx = threadIdx % nSegPerCol;
	uint cIdx = cSegIdx*nSegSize;
	uint cIdx2 = cIdx + nSegSize;
	uint rIdx = rSegIdx;
	uint src_addr1 = rIdx*g_nCol_srcMat + cIdx;
	uint src_addr2 = rIdx*g_nCol_srcMat + cIdx2;
	uint dst_addr = src_addr1;
	float elem1 = 0;
	float elem2 = 0;
	if (cIdx < g_nCol_srcMat && rIdx < g_nRow_srcMat)
	{
		elem1 = g_dstMat[src_addr1];
	}
	if (cIdx2 < g_nCol_srcMat && rIdx < g_nRow_srcMat)
	{
		elem2 = g_dstMat[src_addr2];
	}
	g_dstMat[dst_addr] = elem1 + elem2;
	return;
}


[numthreads( nThread,  1,  1)]    
void cs_Matrix_copyRow(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	if (threadIdx >= g_nCol_dstMat)
		return;
	uint cIdx = threadIdx;
	uint rIdx = g_rowIdx;

	uint addr = rIdx * g_nCol_dstMat + cIdx;
	
	g_dstMat[addr] = g_srcMat[addr];
	return;
}

///////////////////////////////////////////////
uint g_dimVec = 0;
[numthreads( nThread,  1,  1)]    //nThreadGroup = nCol
void cs_VecDot(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	if (threadIdx > g_dimVec)
		return ;
	uint addr = threadIdx;
	g_dstMat[addr] = g_srcMat[addr]*g_src2Mat[addr] ;
	return;
}

[numthreads( nThread,  1,  1)]    //nThreadGroup = nCol
void cs_VecDot_inout(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	if (threadIdx > g_dimVec)
		return ;
	uint addr = threadIdx;
	g_dstMat[addr] = g_dstMat[addr]*g_src2Mat[addr] ;
	return;
}
///////////////////////////////////////////////

[numthreads( nThread,  1,  1)]    //nThreadGroup = nRow*nCol / nThread
void cs_Matrix_Add(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	if (threadIdx > g_dimVec)
		return ;
	uint rIdx = (int)(threadIdx / g_nCol_srcMat);
	uint cIdx = threadIdx % g_nCol_srcMat;
	uint addr = rIdx*g_nCol_srcMat + cIdx;

	g_dstMat[addr] = g_srcMat[addr] + g_src2Mat[addr];
	return;
}

[numthreads( nThread,  1,  1)]    //nThreadGroup = nRow*nCol / nThread
void cs_Matrix_Add_inout(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	if (threadIdx > g_dimVec)
		return ;
	uint rIdx = (int)(threadIdx / g_nCol_dstMat);
	uint cIdx = threadIdx % g_nCol_dstMat;
	uint addr = rIdx*g_nCol_dstMat + cIdx;

	g_dstMat[addr] = g_dstMat[addr] + g_src2Mat[addr];
	return;
}

[numthreads( nThread,  1,  1)]    //nThreadGroup = nRow*nCol / nThread
void cs_Matrix_MulScalar(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	if (threadIdx > g_dimVec)
		return ;
	uint rIdx = (int)(threadIdx / g_nCol_dstMat);
	uint cIdx = threadIdx % g_nCol_dstMat;
	uint addr = rIdx*g_nCol_dstMat + cIdx;

	g_dstMat[addr] = g_srcMat[addr]*g_vScalar;
	return;
}

[numthreads( nThread,  1,  1)]    //nThreadGroup = nRow*nCol / nThread
void cs_Matrix_MulScalar_inout(uint3 Gid : SV_GroupID, uint3 DTid : SV_DispatchThreadID, uint3 GTid : SV_GroupThreadID, uint GI : SV_GroupIndex)
{
	uint threadIdx = (Gid.z*nGroupY*nGroupX + Gid.y*nGroupX + Gid.x)*nThread + GTid.x;
	if (threadIdx > g_dimVec)
		return ;
	uint rIdx = (int)(threadIdx / g_nCol_dstMat);
	uint cIdx = threadIdx % g_nCol_dstMat;
	uint addr = rIdx*g_nCol_dstMat + cIdx;

	g_dstMat[addr] = g_dstMat[addr]*g_vScalar;
	return;
}

technique11 Matrix_GetRowRange {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_GetRowRange()));
    }
}
technique11 Matrix_GetColRange {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_GetColRange()));
    }
}
technique11 Matrix_GetDiag {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_GetDiag()));
    }
}

technique11 Matrix_SetDiag {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_SetDiag()));
    }
}

technique11 Matrix_SetTo {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_SetTo()));
    }
}

technique11 Matrix_MaskSetTo {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_MaskSetTo()));
    }
}


technique11 Matrix_Identity {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_Identity()));
    }
}

technique11 Matrix_MulScalar {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_Matrix_MulScalar()));
    }
}
technique11 Matrix_MulScalar_inout {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_Matrix_MulScalar_inout()));
    }
}

technique11 Matrix_Mul {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_MatrixMul_1RowMul()));
    }
   pass p1 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_MatrixMul_colSum()));
    }
}

technique11 Vec_Dot {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_VecDot()));
    }
}

technique11 Vec_Dot_inout {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_VecDot_inout()));
    }
}

technique11 Matrix_Add {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_Matrix_Add()));
    }
}

technique11 Matrix_Add_inout {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_Matrix_Add_inout()));
    }
}

technique11 Matrix_MaskCopy {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_MaskCopy()));
    }
}

technique11 Matrix_CopyRow {
   pass p0 
	{
        SetVertexShader( NULL );
        SetGeometryShader( NULL );
        SetPixelShader( NULL );
		SetComputeShader( CompileShader( cs_5_0, cs_Matrix_copyRow()));
    }
}
